//
// Copyright 2024 The ANGLE Project Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//

#ifndef COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_
#define COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_

#include "compiler/translator/Common.h"
#include "compiler/translator/Compiler.h"
#include "compiler/translator/IntermNode.h"

namespace sh
{

const char kDefaultUniformBlockVarType[]           = "ANGLE_DefaultUniformBlock";
const char kDefaultUniformBlockVarName[]           = "ANGLE_defaultUniformBlock";
const uint32_t kDefaultUniformBlockBindGroup       = 0;
const uint32_t kDefaultVertexUniformBlockBinding   = 0;
const uint32_t kDefaultFragmentUniformBlockBinding = 1;

const uint32_t kTextureAndSamplerBindGroup = 1;
// The translator emits this dummy location, which needs to be replaced when linking the two
// separate shader stages.
const char kTextureSamplerBindingMarker[] = "@group(1) @binding(@@@@@@) var ";
// The translator emits split samplers/textures for GLSL's combined textures/samplers (combined
// textures/samplers are not supported by WGSL). The new variables are prefixed with these
// constants, respectively.
const char kAngleSamplerPrefix[] = "ANGLE_sampler_";
const char kAngleTexturePrefix[] = "ANGLE_texture_";

const char kWrappedStructFieldName[] = "elem";

struct UniformBlockMetadata
{
    // A list of structs used anywhere in the uniform address space. These will require special
    // handling (@align() attributes, wrapping of basic types, etc.) to ensure they fit WGSL's
    // uniform layout requirements.
    // The key is TSymbolUniqueId::get().
    TUnorderedSet<int> structsInUniformAddressSpace;
};

// Given a GLSL AST `root`, fills in `outMetadata`, to be used when outputting WGSL.
// If the AST is manipulated after calling this, it may be out of sync with the data recorded in
// `outMetadata`.
bool RecordUniformBlockMetadata(TIntermBlock *root, UniformBlockMetadata &outMetadata);

// Based on the GLSL, some extra WGSL will have to be generated so it can be referenced by
// the WGSL generated by the traverser, This tracks exactly which WGSL snippets will need to be
// generated,
struct WGSLGenerationMetadataForUniforms
{
    // Arrays must have a stride of at least 16 if used in the uniform address spaces. If the array
    // element type doesn't have an aligned size of a multiple of 16 (e.g. f32), the element type
    // must be wrapped in a struct which is then aligned to 16. Adding to
    // `arrayElementTypesInUniforms` will cause `OutputUniformWrapperStructsAndConversions` to
    // generate a WGSL wrapper struct of the form:
    //
    // struct ANGLE_wrapper_f32 {
    //   @align(16) elem : f32;
    // };
    TSet<TType> arrayElementTypesInUniforms;

    // If we need to convert arrays with wrapped element types into arrays with unwrapped element
    // types, the necessary conversions are listed here.
    TSet<TType> arrayElementTypesThatNeedUnwrappingConversions;

    // MatCx2 in a uniform will be represented as array<ANGLE_wrapped_vec2, C> to match std140 (WGSL
    // uniforms pack their matrices a bit tighter). These will have to be converted back to matCx2
    // for all other purposes (e.g. multiplication).
    // If this set isn't empty, ANGLE_wrapped_vec2 will be generated even if it
    // hasn't been yet.
    TSet<TType> outputMatCx2Conversion;
};
bool OutputUniformWrapperStructsAndConversions(
    TInfoSinkBase &output,
    const WGSLGenerationMetadataForUniforms &wgslGenerationMetadataForUniforms);

bool IsMatCx2(const TType *type);

ImmutableString MakeUnwrappingArrayConversionFunctionName(const TType *type);
ImmutableString MakeMatCx2ConversionFunctionName(const TType *type);

// TODO(anglebug.com/42267100): for now does not output all uniform blocks,
// just the default block. (fails for  matCx2, bool.)
bool OutputUniformBlocksAndSamplers(TCompiler *compiler, TIntermBlock *root);

// GLSL sampler uniforms are extracted from structs. Given a GLSL sampler's associated name string,
// this function retrieves its new WGSL name and strips off array indices.
std::string WGSLGetMappedSamplerName(const std::string &originalName);

}  // namespace sh

#endif  // COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_
